import io
import pathlib

import PIL.Image
import matplotlib.pyplot
import torch
from lotka_volterra_dataset import LotkaVolterraDataset


def _generate_frame(dataset: LotkaVolterraDataset,
                    left: float,
                    right: float,
                    lim_left: float | None = None,
                    lim_right: float | None = None):
    if lim_left is None:
        lim_left = left
    if lim_right is None:
        lim_right = right

    matplotlib.pyplot.figure(figsize=(8, 5), dpi=110)
    matplotlib.pyplot.xlim(lim_left, lim_right)
    matplotlib.pyplot.ylim(100, 350)

    problem = dataset.problem()
    t = torch.linspace(left, right, 1000)
    x1 = list(problem.x1_batch(t))
    x2 = list(problem.x2_batch(t))
    matplotlib.pyplot.plot(t, x1, label="$x_1$")
    matplotlib.pyplot.plot(t, x2, label="$x_2$")

    t1 = dataset.t1()
    mask1 = (left <= t1) & (t1 < right)
    t1 = t1[mask1]
    x1 = dataset.x1()[mask1]

    t2 = dataset.t2()
    mask2 = (left <= t2) & (t2 < right)
    t2 = t2[mask2]
    x2 = dataset.x2()[mask2]

    matplotlib.pyplot.scatter(t1, x1)
    matplotlib.pyplot.scatter(t2, x2)

    matplotlib.pyplot.legend(loc="lower right")
    matplotlib.pyplot.xlabel(r"$t$")
    matplotlib.pyplot.ylabel(r"$x$")

    stream = io.BytesIO()
    matplotlib.pyplot.savefig(stream)
    matplotlib.pyplot.close()

    stream.seek(0)
    return PIL.Image.open(stream)


def _duration(t: float):
    if t < 15:
        speed = 7 + t ** 3 / 1200
    elif t >= 85:
        if t < 100:
            speed = 13.6 - (t - 85) ** 3 / 1200
        else:
            speed = 7
    else:
        speed = 13.6
    return 600 / speed


def main():
    dataset = LotkaVolterraDataset.load(pathlib.Path(f"./checkpoints/data.pt"))
    dataset.problem().x1(110)
    images = []
    durations = []
    for right in torch.arange(0, 110.5, 0.5):
        lim_right = min(max(right + 2, 10), 110)
        frame = _generate_frame(dataset,
                                left=lim_right - 10,
                                right=right,
                                lim_right=lim_right)
        images.append(frame)
        right = float(right)
        durations.append(_duration(right))
        print(f"{right}: {durations[-1]}")

    images.pop(0).save("./checkpoints/data_generation.gif",
                       save_all=True,
                       append_images=images,
                       duration=durations)
    images[-1].save("./checkpoints/data_generation_final.gif")


if __name__ == "__main__":
    main()
